package com.affinity.impl.spi;

import java.util.BitSet;

import com.affinity.impl.AffinityMaskSpi;
import com.sun.jna.*;
import com.sun.jna.platform.win32.WinDef;
import com.sun.jna.platform.win32.WinDef.DWORD;
import com.sun.jna.ptr.LongByReference;

/**
 * fixme: Class WindowsAffinityMaskSpi is for porn
 *
 * @author cheremin
 * @since 28.04.12,  14:47
 */
public class WindowsAffinityMaskSpi implements AffinityMaskSpi {

    /** @author BegemoT */
    private interface Kernel32Library extends Library {
        public static final Kernel32Library INSTANCE = ( Kernel32Library ) Native.loadLibrary( "kernel32", Kernel32Library.class );

        public int GetProcessAffinityMask( final int pid,
                                           final PointerType lpProcessAffinityMask,
                                           final PointerType lpSystemAffinityMask ) throws LastErrorException;

        public void SetThreadAffinityMask( final int pid,
                                           final WinDef.DWORD lpProcessAffinityMask ) throws LastErrorException;

        public int GetCurrentThread() throws LastErrorException;

        public int GetCurrentProcessorNumber() throws LastErrorException;
    }


    public WindowsAffinityMaskSpi() {
    }

    @Override
    public void pingSpi() {
        final int cpusCount = Runtime.getRuntime().availableProcessors();

        if ( cpusCount > Long.SIZE ) {
            throw new UnsupportedOperationException( cpusCount + " > Long.SIZE -> not supported by current implementation" );
        }

        getCurrentThreadCPU();

        final BitSet affinityMask = getAffinityMask();
        setAffinityMask( affinityMask );
    }

    @Override
    public int getCurrentThreadCPU() {
        try {
            final int cpuNo = syslib().GetCurrentProcessorNumber();
            if ( cpuNo < 0 ) {
                throw new IllegalStateException( "GetCurrentProcessorNumber() return " + cpuNo );
            }
            return cpuNo;
        } catch ( LastErrorException e ) {
            throw new IllegalStateException( "GetCurrentProcessorNumber() errorNo=" + e.getErrorCode(), e );
        }
    }

    private static Kernel32Library syslib() {
        return Kernel32Library.INSTANCE;
    }

    @Override
    public BitSet getAffinityMask() {
        final LongByReference cpuset1 = new LongByReference( 0 );
        final LongByReference cpuset2 = new LongByReference( 0 );
        try {
            final int ret = syslib().GetProcessAffinityMask( -1, cpuset1, cpuset2 );
            if ( ret < 0 ) {
                throw new IllegalStateException( "GetProcessAffinityMask(-1, &(" + cpuset1 + "), &(" + cpuset2 + ") ) return " + ret );
            }
            return Utils.fromLong( cpuset1.getValue() );
        } catch ( LastErrorException e ) {
            throw new IllegalStateException( "GetProcessAffinityMask(-1, &(" + cpuset1 + "), &(" + cpuset2 + ") ) errorNo=" + e.getErrorCode(), e );
        }
    }

    @Override
    public void setAffinityMask( final BitSet affinityMask ) {
        if ( affinityMask == null || affinityMask.isEmpty() ) {
            throw new IllegalArgumentException( "affinityMask can't be null nor empty" );
        }
        final DWORD aff = new DWORD( Utils.toLong( affinityMask ) );
        final int tid;
        try {
            tid = syslib().GetCurrentThread();
        } catch ( LastErrorException e ) {
            throw new IllegalStateException( "GetCurrentThread() errorNo=" + e.getErrorCode(), e );
        }
        try {
            syslib().SetThreadAffinityMask( tid, aff );
        } catch ( LastErrorException e ) {
            throw new IllegalStateException( "SetThreadAffinityMask((" + tid + ") , &(" + affinityMask + ") ) errorNo=" + e.getErrorCode(), e );
        }
    }
}
